package top.mcwebsite.lib.router.plugin

import org.objectweb.asm.ClassReader
import org.objectweb.asm.ClassWriter
import org.objectweb.asm.tree.ClassNode
import top.mcwebsite.lib.transform.api.TransformContext
import top.mcwebsite.lib.transform.api.Transformer

class RouterTransformer : Transformer {

    companion object {

        const val ROUTER_BUILDER_INJECT_CLASS = "top/mcwebsite/lib/router/compose/RouteGraphBuilder"
    }

    private val routerBuilderModuleNames = mutableSetOf<String>()

    override fun isNeedScan(): Boolean {
        return true
    }

    override fun scan(path: String, byteArraySupplier: () -> ByteArray) {
        if (!path.contains(Constant.ROUTER_BUILDER_PREFIX)) {
            return
        }
        if (path.contains("$")) {
            return
        }
        val start = path.indexOf(Constant.ROUTER_BUILDER_PREFIX)
        val end = path.length - "Kt.class".length
        val funcName = path.substring(start + Constant.ROUTER_BUILDER_PREFIX.length, end)
        routerBuilderModuleNames.add(funcName)
    }

    override fun onBeforeTransform(context: TransformContext) {
    }

    override fun transform(byteArray: ByteArray): ByteArray {
        val reader = ClassReader(byteArray)
        val classNode = ClassNode()
        reader.accept(classNode, 0)
        return if (classNode.name.contains(ROUTER_BUILDER_INJECT_CLASS)) {
            println("router builder = $routerBuilderModuleNames")
            val writer = ClassWriter(reader, 0)
            val visitor = RouteBuilderCodeClassVisitor(writer, routerBuilderModuleNames)
            reader.accept(visitor, 0)
            writer.toByteArray()
        } else {
            byteArray
        }
    }

    override fun onAfterTransform(context: TransformContext) {
    }
}